﻿using Microsoft.Extensions.AI;
using ModelContextProtocol.Client;
using OpenAI;
using System.ClientModel;

namespace MCPSharp.Client;

/// <summary>
/// Represents a wrapper client class for interacting with AI chat models.
/// Responsible for initializing the chat client and maintaining conversation context.
/// </summary>
public class ChatAIClient
{
    /// <summary>
    /// The encapsulated AI chat client interface, supporting features such as function calling.
    /// </summary>
    private IChatClient ChatClient;

    /// <summary>
    /// Stores all chat message records of the current session.
    /// </summary>
    private IList<ChatMessage> Messages;

    /// <summary>
    /// API access key for authentication. [REPLACE WITH YOUR OWN KEY]
    /// </summary>
    private const string _apiKey = "sk-xxxx";

    /// <summary>
    /// Base request URL for the AI service. [REPLACE WITH YOUR OWN ENDPOINT]
    /// </summary>
    private const string _baseURL = "https://api.siliconflow.cn/v1";

    /// <summary>
    /// Identifier of the AI model to use. [REPLACE WITH YOUR OWN MODEL ID]
    /// </summary>
    private const string _modelID = "moonshotai/Kimi-K2-Instruct";

    /// <summary>
    /// Initializes a new instance of the <see cref="ChatAIClient"/> class.
    /// Automatically configures the chat client during construction.
    /// </summary>
    public ChatAIClient()
    {
        InitIChatClient();
    }

    /// <summary>
    /// Initializes the internal AI chat client instance.
    /// Configures API credentials, service endpoint, and builds a client capable of function calling.
    /// Also initializes the system message as the starting point of the conversation.
    /// </summary>
    private void InitIChatClient()
    {
        // Create API key credential
        ApiKeyCredential apiKeyCredential = new ApiKeyCredential(_apiKey);

        // Configure OpenAI client options, e.g., custom endpoint
        OpenAIClientOptions openAIClientOptions = new OpenAIClientOptions();
        openAIClientOptions.Endpoint = new Uri(_baseURL);

        // Create OpenAI client and obtain the chat interface for the specified model
        var openaiClient = new OpenAIClient(apiKeyCredential, openAIClientOptions)
            .GetChatClient(_modelID)
            .AsIChatClient();

        // Build an enhanced chat client (e.g., enable function invocation)
        ChatClient = new ChatClientBuilder(openaiClient)
            .UseFunctionInvocation()
            .Build();

        // Initialize conversation history with a system prompt message
        Messages =
        [
            // Add system role message
            new(ChatRole.System, "You are a helpful assistant helping us test MCP server functionality. Prefer answering in Chinese!"),
        ];
    }

    /// <summary>
    /// Asynchronously processes the user's natural language query and interacts with the AI model, supporting MCP tool invocations.
    /// </summary>
    /// <param name="query">The user's natural language input</param>
    /// <param name="tools">List of available MCP tools to extend the AI's external capabilities</param>
    /// <returns>The final text response generated by the AI</returns>
    public async Task<string> ProcessQueryAsync(string query, IList<McpClientTool> tools)
    {
        // If message history is empty, initialize with system prompt
        if (Messages.Count == 0)
        {
            Messages =
            [
                new(ChatRole.System, "You are a helpful assistant helping us test MCP server functionality. Prefer answering in Chinese!")
            ];
        }

        // Add user input message to conversation history
        Messages.Add(new(ChatRole.User, query));

        // Set request options and inject available tools
        var options = new ChatOptions
        {
            Tools = [.. tools]
        };

        // Call AI client to get streaming response
        Console.ForegroundColor = ConsoleColor.Green;

        List<ChatResponseUpdate> updates = [];
        await foreach (ChatResponseUpdate update in ChatClient.GetStreamingResponseAsync(Messages, options))
        {
            Task task = new Task(() =>
            {
                if (!string.IsNullOrEmpty(update.Text))
                    Console.Write(update);
            });
            task.Start();
            updates.Add(update);
        }
        Console.ForegroundColor = ConsoleColor.Yellow;

        //var response = await ChatClient.GetResponseAsync(Messages, options);

        // Add AI response to conversation history
        //Messages.AddMessages(response);

        //// Output information about invoked tools
        //OutputToolUsageInfo(response);

        // Return empty string for now (placeholder)
        return "";
    }

    /// <summary>
    /// Helper method: Outputs tool usage information invoked by the AI during response to the console.
    /// </summary>
    /// <param name="response">The full response object from the AI</param>
    private void OutputToolUsageInfo(ChatResponse response)
    {
        // Get all messages with Tool role
        var toolUseMessages = response.Messages.Where(m => m.Role == ChatRole.Tool).ToList();

        // Determine if any tools were invoked
        var toolUseMessage = response.Messages.Where(m => m.Role == ChatRole.Tool);

        // Check if the first message contains more than one content item (typically: user question + function call)
        if (response.Messages[0].Contents.Count > 1)
        {
            // Try to extract function call info from the second content item of the first message
            var functionCall = (FunctionCallContent)response.Messages[0].Contents[1];

            // Set console color to green to highlight tool usage info
            Console.ForegroundColor = ConsoleColor.Green;

            string arguments = "";

            // If function call includes arguments, concatenate them
            if (functionCall.Arguments != null)
            {
                foreach (var arg in functionCall.Arguments)
                {
                    arguments += $"{arg.Key}:{arg.Value};";
                }

                // Output called function name and its arguments
                Console.WriteLine($"Invoked Function: {functionCall.Name}; Arguments: {arguments}");

                // Iterate through all Tool messages and output each tool execution result
                foreach (var message in toolUseMessage)
                {
                    // Extract execution result after tool invocation
                    var functionResultContent = (FunctionResultContent)message.Contents[0];
                    Console.WriteLine($"Tool Execution Result: {functionResultContent.Result}");
                }

                // Restore console default color (white)
                Console.ForegroundColor = ConsoleColor.White;
            }
            else
            {
                // No arguments provided
                Console.WriteLine("Tool arguments are empty.");
            }
        }
        else
        {
            Console.ForegroundColor = ConsoleColor.Green;
            Console.WriteLine("No tools were invoked this time.");
            Console.ForegroundColor = ConsoleColor.White;
        }
    }
}